#!/usr/bin/env python
# -*- coding: utf-8 -*-

import threading
import time
import logging
from pathlib import Path
from typing import List, Tuple, Optional, Dict, Any
from dataclasses import dataclass

from .base_converter import BaseConverter, FileInfo

logger = logging.getLogger(__name__)

@dataclass
class OCRInstance:
    """OCR实例信息"""
    instance: Any
    created_time: float
    last_used_time: float
    use_count: int = 0

class OCRInstanceManager:
    """OCR实例管理器 - 单例模式管理OCR模型实例"""
    
    _instance = None
    _lock = threading.Lock()
    
    def __new__(cls):
        if cls._instance is None:
            with cls._lock:
                if cls._instance is None:
                    cls._instance = super().__new__(cls)
                    cls._instance._initialized = False
        return cls._instance
    
    def __init__(self):
        if self._initialized:
            return
            
        self._initialized = True
        self._ocr_instances: Dict[str, OCRInstance] = {}
        self._instance_locks: Dict[str, threading.Lock] = {}
        self._global_lock = threading.Lock()
        
        logger.info("OCR实例管理器已初始化")
    
    def get_ocr_instance(self, ocr_type: str) -> Any:
        """获取OCR实例（单例模式）"""
        if ocr_type in self._ocr_instances:
            instance = self._ocr_instances[ocr_type]
            instance.last_used_time = time.time()
            instance.use_count += 1
            logger.debug(f"重用OCR实例: {ocr_type} (使用次数: {instance.use_count})")
            return instance.instance
        
        # 创建实例锁（如果不存在）
        if ocr_type not in self._instance_locks:
            with self._global_lock:
                if ocr_type not in self._instance_locks:
                    self._instance_locks[ocr_type] = threading.Lock()
        
        # 线程安全创建实例
        with self._instance_locks[ocr_type]:
            # 双重检查
            if ocr_type in self._ocr_instances:
                instance = self._ocr_instances[ocr_type]
                instance.last_used_time = time.time()
                instance.use_count += 1
                return instance.instance
            
            # 创建新的OCR实例
            logger.info(f"创建新的OCR实例: {ocr_type}")
            start_time = time.time()
            
            try:
                ocr_instance = self._create_ocr_instance(ocr_type)
                
                # 保存实例信息
                instance_info = OCRInstance(
                    instance=ocr_instance,
                    created_time=start_time,
                    last_used_time=time.time(),
                    use_count=1
                )
                
                self._ocr_instances[ocr_type] = instance_info
                
                end_time = time.time()
                logger.info(f"OCR实例创建完成: {ocr_type}, 耗时: {end_time - start_time:.3f}秒")
                
                return ocr_instance
                
            except Exception as e:
                logger.error(f"创建OCR实例失败: {ocr_type}, 错误: {e}")
                raise
    
    def _create_ocr_instance(self, ocr_type: str) -> Any:
        """创建OCR实例"""
        if ocr_type == 'paddleocr':
            from paddleocr import PaddleOCR
            # 参照ocrToMd.py的配置，优化文档识别效果
            return PaddleOCR(
                use_doc_orientation_classify=False,
                use_doc_unwarping=False,
                use_textline_orientation=False,
                lang='ch',
                show_log=False
            )
        elif ocr_type == 'easyocr':
            import easyocr
            return easyocr.Reader(['ch_sim', 'en'])
        elif ocr_type == 'tesseract':
            import pytesseract
            # Tesseract 不需要实例，返回模块本身
            return pytesseract
        else:
            raise ValueError(f"不支持的OCR类型: {ocr_type}")
    
    def cleanup_instance(self, ocr_type: str):
        """清理OCR实例"""
        if ocr_type in self._ocr_instances:
            with self._instance_locks.get(ocr_type, threading.Lock()):
                if ocr_type in self._ocr_instances:
                    del self._ocr_instances[ocr_type]
                    logger.info(f"已清理OCR实例: {ocr_type}")
    
    def cleanup_all_instances(self):
        """清理所有OCR实例"""
        with self._global_lock:
            for ocr_type in list(self._ocr_instances.keys()):
                self.cleanup_instance(ocr_type)
        logger.info("所有OCR实例已清理")
    
    def get_statistics(self) -> Dict[str, Any]:
        """获取OCR实例统计信息"""
        stats = {
            'total_instances': len(self._ocr_instances),
            'instances': {},
            'total_usage': 0
        }
        
        for ocr_type, instance in self._ocr_instances.items():
            stats['instances'][ocr_type] = {
                'use_count': instance.use_count,
                'created_time': instance.created_time,
                'last_used_time': instance.last_used_time,
                'uptime_hours': (time.time() - instance.created_time) / 3600
            }
            stats['total_usage'] += instance.use_count
        
        return stats

# 全局OCR实例管理器
ocr_manager = OCRInstanceManager()

class ImageConverter(BaseConverter):
    """图片转换器"""
    
    OPTIONAL_DEPENDENCIES = {
        'paddleocr': 'paddleocr',
        'easyocr': 'easyocr',
        'tesseract': 'pytesseract'
    }
    
    def __init__(self, temp_dir: Path, file_id: Optional[str] = None):
        super().__init__(temp_dir, file_id)
        self.available_ocr_engines = self._check_ocr_engines()
        logger.info(f"图片转换器初始化完成，可用OCR引擎: {list(self.available_ocr_engines.keys())}")
    
    def _check_ocr_engines(self) -> Dict[str, bool]:
        """检查可用的OCR引擎"""
        engines = {}
        
        # 检查 PaddleOCR
        try:
            import paddleocr
            engines['paddleocr'] = True
            logger.debug("PaddleOCR 可用")
        except ImportError:
            engines['paddleocr'] = False
            logger.debug("PaddleOCR 不可用")
        
        # 检查 EasyOCR
        try:
            import easyocr
            engines['easyocr'] = True
            logger.debug("EasyOCR 可用")
        except ImportError:
            engines['easyocr'] = False
            logger.debug("EasyOCR 不可用")
        
        # 检查 Tesseract
        try:
            import pytesseract
            # 尝试运行tesseract命令来确认安装
            pytesseract.get_tesseract_version()
            engines['tesseract'] = True
            logger.debug("Tesseract 可用")
        except ImportError:
            engines['tesseract'] = False
            logger.debug("Tesseract 不可用 (pytesseract未安装)")
        except Exception as e:
            engines['tesseract'] = False
            logger.debug(f"Tesseract 不可用: {e}")
        
        return engines
    
    def get_available_methods(self) -> List[str]:
        """获取可用的转换方法"""
        methods = []
        for engine, available in self.available_ocr_engines.items():
            if available:
                methods.append(engine)
        return methods
    
    def is_dependency_available(self, dependency: str) -> bool:
        """检查依赖是否可用"""
        return self.available_ocr_engines.get(dependency, False)
    
    async def convert(self, file_path: Path, method: Optional[str] = None, **kwargs) -> Tuple[str, List[FileInfo]]:
        """
        转换图片文件为Markdown
        
        Args:
            file_path: 图片文件路径
            method: OCR方法 ('paddleocr', 'easyocr', 'tesseract')
            **kwargs: 其他参数
                - y_tolerance: Y坐标容差 (默认: 10)
                - paragraph_gap: 段落间距阈值 (默认: 30)
                - confidence_threshold: 置信度阈值 (默认: 0.5)
                - enable_layout_analysis: 是否启用版面分析 (默认: True)
            
        Returns:
            tuple: (markdown内容, 文件信息列表)
        """
        available_methods = self.get_available_methods()
        if not available_methods:
            raise ImportError(
                "没有可用的OCR引擎。请安装以下任一包：\n"
                "- pip install paddleocr\n"
                "- pip install easyocr\n"
                "- pip install pytesseract (需要系统安装tesseract)"
            )
        
        # 选择OCR方法
        if method and method not in available_methods:
            raise ImportError(f"OCR方法 '{method}' 不可用。可用方法: {', '.join(available_methods)}")
        
        if not method:
            # 按优先级选择：paddleocr > easyocr > tesseract
            if 'paddleocr' in available_methods:
                method = 'paddleocr'
            elif 'easyocr' in available_methods:
                method = 'easyocr'
            else:
                method = 'tesseract'
        
        logger.info(f"使用 {method} 进行OCR识别: {file_path.name}")
        
        try:
            # 根据不同的OCR引擎进行处理
            if method == 'paddleocr':
                markdown_content = await self._convert_with_paddleocr(file_path, **kwargs)
            elif method == 'easyocr':
                markdown_content = await self._convert_with_easyocr(file_path, **kwargs)
            elif method == 'tesseract':
                markdown_content = await self._convert_with_tesseract(file_path, **kwargs)
            else:
                raise ValueError(f"不支持的OCR方法: {method}")
            
            # 创建文件信息（图片转换通常不产生额外文件）
            file_info = FileInfo(
                file_path=file_path,
                file_id=file_path.stem,
                file_name=file_path.name,
                file_size=file_path.stat().st_size,
                file_type="image",
                mime_type=self._get_mime_type(file_path)
            )
            
            logger.info(f"图片转换完成: {file_path.name}")
            return markdown_content, [file_info]
            
        except Exception as e:
            logger.error(f"图片转换失败: {file_path.name}, 错误: {e}")
            raise
    
    async def _convert_with_paddleocr(self, file_path: Path, **kwargs) -> str:
        """使用PaddleOCR进行转换 - 参照ocrToMd.py的高级文本排版算法"""
        # 获取OCR实例（单例模式）
        ocr_instance = ocr_manager.get_ocr_instance('paddleocr')
        
        # 执行OCR识别 - 使用predict方法以获得更好的结果格式
        try:
            result = ocr_instance.predict(input=str(file_path))
            if not result or len(result) == 0:
                return f"# 图片OCR结果\n\n**文件名：** {file_path.name}\n\n*未识别到文本内容*\n"
            
            # 使用ocrToMd.py的文本排版算法
            content = self._ocr_result_to_markdown(
                result[0], 
                y_tolerance=kwargs.get('y_tolerance', 10),
                paragraph_gap=kwargs.get('paragraph_gap', 30)
            )
            
        except AttributeError:
            # 如果predict方法不可用，回退到ocr方法
            logger.warning("predict方法不可用，使用ocr方法")
            result = ocr_instance.ocr(str(file_path), cls=True)
            
            if not result or not result[0]:
                return f"# 图片OCR结果\n\n**文件名：** {file_path.name}\n\n*未识别到文本内容*\n"
            
            # 简单文本提取（回退方案）
            content = self._simple_ocr_result_to_text(
                result[0], 
                confidence_threshold=kwargs.get('confidence_threshold', 0.5)
            )
        
        # 构建markdown格式
        markdown_lines = [
            f"# 图片OCR结果",
            f"",
            f"**文件名：** {file_path.name}",
            f"**识别引擎：** PaddleOCR",
            f"",
            f"## 识别内容",
            f"",
            content.strip()
        ]
        
        return "\n".join(markdown_lines)
    
    def _ocr_result_to_markdown(self, result: dict, y_tolerance: int = 10, paragraph_gap: int = 30) -> str:
        """
        将OCR结果转换为markdown - 参照ocrToMd.py的智能排版算法
        
        Args:
            result: OCR结果字典，包含rec_texts和rec_boxes
            y_tolerance: Y坐标容差
            paragraph_gap: 段落间距阈值
        """
        if "rec_texts" not in result or "rec_boxes" not in result:
            return "*OCR结果格式不正确*"
        
        texts = result["rec_texts"]
        boxes = result["rec_boxes"]
        
        if not texts or not boxes:
            return "*未识别到文本内容*"
        
        # 构建包含位置信息的文本行
        lines = []
        for text, box in zip(texts, boxes):
            if isinstance(box, (list, tuple)) and len(box) >= 4:
                x1, y1, x2, y2 = box[:4]
                top_y = y1
                left_x = x1
                lines.append({"text": text, "top_y": top_y, "left_x": left_x})
        
        if not lines:
            return "*未能解析OCR结果*"
        
        # 按垂直位置排序，然后按水平位置排序
        lines.sort(key=lambda l: (l["top_y"], l["left_x"]))
        
        # 智能合并同行和分段落
        paragraphs = []
        current_line = []
        prev_y = -1000
        
        for line in lines:
            # 如果Y坐标差距大于段落间距阈值，开始新段落
            if abs(line["top_y"] - prev_y) > paragraph_gap:
                if current_line:
                    # 将当前行按X坐标排序并合并
                    merged_line = " ".join([l["text"] for l in sorted(current_line, key=lambda l: l["left_x"])])
                    if merged_line.strip():
                        paragraphs.append(merged_line.strip())
                current_line = [line]
            else:
                # 继续当前行
                current_line.append(line)
            prev_y = line["top_y"]
        
        # 处理最后一行
        if current_line:
            merged_line = " ".join([l["text"] for l in sorted(current_line, key=lambda l: l["left_x"])])
            if merged_line.strip():
                paragraphs.append(merged_line.strip())
        
        # 用双换行符连接段落
        return "\n\n".join(paragraphs)
    
    def _simple_ocr_result_to_text(self, result_data: list, confidence_threshold: float = 0.5) -> str:
        """简单的OCR结果转文本（回退方案）"""
        text_lines = []
        
        for line_info in result_data:
            if line_info and len(line_info) >= 2:
                # line_info: [[[x1,y1],[x2,y2],[x3,y3],[x4,y4]], (text, confidence)]
                text_info = line_info[1]  # (文本, 置信度)
                
                if isinstance(text_info, (list, tuple)) and len(text_info) >= 2:
                    text, confidence = text_info
                    # 使用可配置的置信度阈值
                    if confidence > confidence_threshold:
                        text_lines.append(text.strip())
        
        return "\n".join(text_lines)
    
    async def _convert_with_easyocr(self, file_path: Path, **kwargs) -> str:
        """使用EasyOCR进行转换"""
        # 获取OCR实例（单例模式）
        reader = ocr_manager.get_ocr_instance('easyocr')
        
        # 执行OCR识别
        results = reader.readtext(str(file_path))
        
        if not results:
            return f"# 图片OCR结果\n\n**文件名：** {file_path.name}\n\n*未识别到文本内容*\n"
        
        # 处理OCR结果
        markdown_lines = [
            f"# 图片OCR结果",
            f"",
            f"**文件名：** {file_path.name}",
            f"**识别引擎：** EasyOCR",
            f"",
            f"## 识别内容",
            f""
        ]
        
        for result in results:
            # result: (box, text, confidence)
            if len(result) >= 3:
                box, text, confidence = result
                # 只包含置信度较高的文本
                if confidence > 0.5:
                    markdown_lines.append(text.strip())
        
        return "\n".join(markdown_lines)
    
    async def _convert_with_tesseract(self, file_path: Path, **kwargs) -> str:
        """使用Tesseract进行转换"""
        # 获取OCR实例（实际是pytesseract模块）
        pytesseract = ocr_manager.get_ocr_instance('tesseract')
        
        # 直接使用PIL读取图片（避免OpenCV依赖）
        try:
            from PIL import Image
            image = Image.open(str(file_path))
            
            # 执行OCR识别
            try:
                text = pytesseract.image_to_string(image, lang='chi_sim+eng')
            except Exception as e:
                # 如果中文识别失败，尝试只用英文
                logger.warning(f"中文OCR失败，尝试英文识别: {e}")
                text = pytesseract.image_to_string(image, lang='eng')
            
        except ImportError:
            raise ImportError("Tesseract转换需要PIL库: pip install Pillow")
        except Exception as e:
            raise ValueError(f"无法读取图片文件: {file_path}, 错误: {e}")
        
        if not text.strip():
            return f"# 图片OCR结果\n\n**文件名：** {file_path.name}\n\n*未识别到文本内容*\n"
        
        # 处理OCR结果
        markdown_lines = [
            f"# 图片OCR结果",
            f"",
            f"**文件名：** {file_path.name}",
            f"**识别引擎：** Tesseract",
            f"",
            f"## 识别内容",
            f"",
            text.strip()
        ]
        
        return "\n".join(markdown_lines)
    
    def _get_mime_type(self, file_path: Path) -> str:
        """获取文件的MIME类型"""
        ext = file_path.suffix.lower()
        mime_types = {
            '.png': 'image/png',
            '.jpg': 'image/jpeg',
            '.jpeg': 'image/jpeg',
            '.gif': 'image/gif',
            '.bmp': 'image/bmp',
            '.tiff': 'image/tiff',
            '.webp': 'image/webp'
        }
        return mime_types.get(ext, 'image/unknown')
    
    def cleanup(self):
        """清理资源"""
        # 转换器本身不需要特殊清理，OCR实例由管理器统一管理
        pass
    
    def get_ocr_statistics(self) -> Dict[str, Any]:
        """获取OCR使用统计信息"""
        return ocr_manager.get_statistics() 